import pygame
import numpy as np
from pygame.locals import *

# 初始化配置
pygame.init()
SCREEN_WIDTH = 1400
SCREEN_HEIGHT = 800
CUBE_SIZE = 250
FRAME_COLOR = (100, 100, 120)
BG_COLOR = (20, 20, 30)

class Math4D:
    """四维数学运算工具类"""
    
    @staticmethod
    def rotation_matrix_4d(xy, xz, xw, yz, yw, zw):
        """生成四维旋转矩阵"""
        matrices = [
            Math4D._rotation_xy(xy),
            Math4D._rotation_xz(xz),
            Math4D._rotation_xw(xw),
            Math4D._rotation_yz(yz),
            Math4D._rotation_yw(yw),
            Math4D._rotation_zw(zw)
        ]
        result = np.eye(4)
        for mat in matrices:
            result = result @ mat
        return result
    
    @staticmethod
    def _rotation_xy(angle):
        c, s = np.cos(angle), np.sin(angle)
        return np.array([
            [c, -s, 0, 0],
            [s, c, 0, 0],
            [0, 0, 1, 0],
            [0, 0, 0, 1]
        ])
    
    @staticmethod
    def _rotation_xz(angle):
        c, s = np.cos(angle), np.sin(angle)
        return np.array([
            [c, 0, -s, 0],
            [0, 1, 0, 0],
            [s, 0, c, 0],
            [0, 0, 0, 1]
        ])
    
    @staticmethod
    def _rotation_xw(angle):
        c, s = np.cos(angle), np.sin(angle)
        return np.array([
            [c, 0, 0, -s],
            [0, 1, 0, 0],
            [0, 0, 1, 0],
            [s, 0, 0, c]
        ])
    
    @staticmethod
    def _rotation_yz(angle):
        c, s = np.cos(angle), np.sin(angle)
        return np.array([
            [1, 0, 0, 0],
            [0, c, -s, 0],
            [0, s, c, 0],
            [0, 0, 0, 1]
        ])
    
    @staticmethod
    def _rotation_yw(angle):
        c, s = np.cos(angle), np.sin(angle)
        return np.array([
            [1, 0, 0, 0],
            [0, c, 0, -s],
            [0, 0, 1, 0],
            [0, s, 0, c]
        ])
    
    @staticmethod
    def _rotation_zw(angle):
        c, s = np.cos(angle), np.sin(angle)
        return np.array([
            [1, 0, 0, 0],
            [0, 1, 0, 0],
            [0, 0, c, -s],
            [0, 0, s, c]
        ])

class Tesseract:
    """四维超立方体类"""
    
    def __init__(self):
        self.vertices_4d = self._create_vertices()
        self.edges = self._create_edges()
        self.faces = self._create_faces()
        self.rotation = np.eye(4)
        self.position_4d = np.array([0.0, 0.0, 0.0, 0.0])
    
    def _create_vertices(self):
        """创建16个四维顶点"""
        vertices = []
        for i in range(16):
            x = 1 if i & 1 else -1
            y = 1 if i & 2 else -1
            z = 1 if i & 4 else -1
            w = 1 if i & 8 else -1
            vertices.append(np.array([x, y, z, w]))
        return vertices
    
    def _create_edges(self):
        """创建32条边"""
        edges = []
        for i in range(16):
            for j in range(i + 1, 16):
                if bin(i ^ j).count('1') == 1:
                    edges.append((i, j))
        return edges
    
    def _create_faces(self):
        """创建24个面"""
        faces = []
        for z in [0, 4]:
            for w in [0, 8]:
                base = z + w
                faces.append([base, base + 1, base + 3, base + 2])
        for y in [0, 2]:
            for w in [0, 8]:
                base = y + w
                faces.append([base, base + 1, base + 5, base + 4])
        for y in [0, 2]:
            for z in [0, 4]:
                base = y + z
                faces.append([base, base + 1, base + 9, base + 8])
        for x in [0, 1]:
            for w in [0, 8]:
                base = x + w
                faces.append([base, base + 2, base + 6, base + 4])
        for x in [0, 1]:
            for z in [0, 4]:
                base = x + z
                faces.append([base, base + 2, base + 10, base + 8])
        for x in [0, 1]:
            for y in [0, 2]:
                base = x + y
                faces.append([base, base + 4, base + 12, base + 8])
        return faces
    
    def update_rotation(self, angles):
        """更新旋转矩阵"""
        self.rotation = Math4D.rotation_matrix_4d(*angles)
    
    def get_rotated_vertices(self):
        """获取旋转和平移后的顶点"""
        return [self.rotation @ v + self.position_4d for v in self.vertices_4d]
    
    def get_3d_slice(self, w_value):
        """获取在特定W值处的三维切面"""
        rotated = self.get_rotated_vertices()
        slice_vertices = []
        slice_edges = []
        vertex_map = {}
        
        for edge_idx, (i, j) in enumerate(self.edges):
            v1, v2 = rotated[i], rotated[j]
            w1, w2 = v1[3], v2[3]
            
            if (w1 - w_value) * (w2 - w_value) <= 0:
                if abs(w2 - w1) < 1e-6:
                    t = 0.5
                else:
                    t = (w_value - w1) / (w2 - w1)
                intersection = v1 + t * (v2 - v1)
                vertex_idx = len(slice_vertices)
                slice_vertices.append(intersection[:3])
                vertex_map[edge_idx] = vertex_idx
        
        for face in self.faces:
            intersecting_edges = []
            for k in range(4):
                v1_idx, v2_idx = face[k], face[(k + 1) % 4]
                for edge_idx, (i, j) in enumerate(self.edges):
                    if (i == v1_idx and j == v2_idx) or (i == v2_idx and j == v1_idx):
                        if edge_idx in vertex_map:
                            intersecting_edges.append(edge_idx)
                        break
            
            if len(intersecting_edges) >= 2:
                for k in range(len(intersecting_edges)):
                    e1 = intersecting_edges[k]
                    e2 = intersecting_edges[(k + 1) % len(intersecting_edges)]
                    if vertex_map[e1] != vertex_map[e2]:
                        slice_edges.append((vertex_map[e1], vertex_map[e2]))
        
        return slice_vertices, slice_edges

class ObservationCube:
    """三维观察立方体"""
    
    def __init__(self, w_value, position, size=CUBE_SIZE, slice_scale=1.0):
        self.w_value = w_value
        self.position = position
        self.size = size
        self.slice_scale = slice_scale
        self.view_rotation_x = 0.4
        self.view_rotation_y = 0.6
        self.frame_size = 1.5
        self.frame_vertices = self._create_frame_vertices()
        self.frame_edges = self._create_frame_edges()
    
    def _create_frame_vertices(self):
        """创建观察立方体框架的8个顶点"""
        s = self.frame_size
        return [
            np.array([-s, -s, -s]), np.array([s, -s, -s]),
            np.array([s, s, -s]), np.array([-s, s, -s]),
            np.array([-s, -s, s]), np.array([s, -s, s]),
            np.array([s, s, s]), np.array([-s, s, s])
        ]
    
    def _create_frame_edges(self):
        """创建框架的12条边"""
        return [
            (0, 1), (1, 2), (2, 3), (3, 0),
            (4, 5), (5, 6), (6, 7), (7, 4),
            (0, 4), (1, 5), (2, 6), (3, 7)
        ]
    
    def contains_point(self, pos):
        """检查点是否在立方体区域内"""
        x, y = pos
        px, py = self.position
        return px <= x <= px + self.size and py <= y <= py + self.size
    
    def project_3d_to_2d(self, vertex_3d):
        """将三维坐标投影到二维屏幕"""
        rx = self.view_rotation_x
        ry = self.view_rotation_y
        cos_rx, sin_rx = np.cos(rx), np.sin(rx)
        cos_ry, sin_ry = np.cos(ry), np.sin(ry)
        
        y = vertex_3d[1] * cos_rx - vertex_3d[2] * sin_rx
        z = vertex_3d[1] * sin_rx + vertex_3d[2] * cos_rx
        x = vertex_3d[0] * cos_ry + z * sin_ry
        z = -vertex_3d[0] * sin_ry + z * cos_ry
        
        distance = 8.0
        scale = distance / (distance + z + 0.1)
        screen_x = int(self.size / 2 + x * self.size * 0.25 * scale)
        screen_y = int(self.size / 2 - y * self.size * 0.25 * scale)
        return (screen_x, screen_y), z
    
    def render(self, surface, slice_vertices, slice_edges, highlight=False):
        """渲染观察立方体"""
        surface.fill((0, 0, 0, 0))
        self._draw_frame(surface, highlight)
        self._draw_slice(surface, slice_vertices, slice_edges, highlight)
        self._draw_label(surface, highlight)
    
    def _draw_frame(self, surface, highlight):
        """绘制立方体框架"""
        color = (120, 120, 140) if highlight else FRAME_COLOR
        width = 2 if highlight else 1
        projected = []
        depths = []
        
        for v in self.frame_vertices:
            pos_2d, depth = self.project_3d_to_2d(v)
            projected.append(pos_2d)
            depths.append(depth)
        
        edge_depth = []
        for i, j in self.frame_edges:
            depth = (depths[i] + depths[j]) / 2
            edge_depth.append((depth, i, j))
        edge_depth.sort()
        
        for _, i, j in edge_depth:
            pygame.draw.line(surface, color, projected[i], projected[j], width)
    
    def _draw_slice(self, surface, vertices, edges, highlight):
        """绘制切面"""
        if not vertices:
            return
        
        scaled_vertices = [v * self.slice_scale for v in vertices]
        projected = []
        depths = []
        
        for v in scaled_vertices:
            pos_2d, depth = self.project_3d_to_2d(v)
            projected.append(pos_2d)
            depths.append(depth)
        
        if len(projected) >= 3:
            try:
                color = (100, 150, 255, 150) if highlight else (80, 120, 200, 100)
                temp_surface = pygame.Surface((self.size, self.size), pygame.SRCALPHA)
                pygame.draw.polygon(temp_surface, color, projected)
                surface.blit(temp_surface, (0, 0))
            except:
                pass
        
        edge_color = (150, 200, 255) if highlight else (100, 150, 200)
        edge_width = 3 if highlight else 2
        drawn_edges = set()
        
        for i, j in edges:
            if i < len(projected) and j < len(projected):
                edge_key = tuple(sorted([i, j]))
                if edge_key not in drawn_edges:
                    pygame.draw.line(surface, edge_color, projected[i], projected[j], edge_width)
                    drawn_edges.add(edge_key)
        
        vertex_color = (200, 230, 255) if highlight else (150, 180, 220)
        vertex_radius = 5 if highlight else 3
        for pos in projected:
            pygame.draw.circle(surface, vertex_color, pos, vertex_radius)
    
    def _draw_label(self, surface, highlight):
        """绘制W值标签"""
        font = pygame.font.Font(None, 20)
        color = (255, 255, 255) if highlight else (180, 180, 180)
        text = font.render(f"W={self.w_value:.2f}", True, color)
        surface.blit(text, (5, 5))

class ControlPanel:
    """控制面板"""
    
    def __init__(self):
        self.font_title = pygame.font.Font(None, 24)
        self.font_normal = pygame.font.Font(None, 16)
        self.rotation_speed = 0.5
        self.translation_speed = 0.05
        self.slice_scale = 1.0
    
    def render(self, surface, tesseract, auto_rotate, paused, auto_move):
        """渲染控制面板"""
        x, y = 10, 10
        
        # 标题
        title = self.font_title.render("4D控制", True, (255, 255, 255))
        surface.blit(title, (x, y))
        y += 30
        
        # 控制说明
        controls = [
            "空格:暂停/继续",
            f"自动旋转:{'开' if auto_rotate else '关'} (R切换)",
            f"自动移动:{'开' if auto_move else '关'} (Tab切换)",
            "",
            "=== 旋转(1-6键) ===",
            "1/Shift+1:XY平面",
            "2/Shift+2:XZ平面",
            "3/Shift+3:XW平面",
            "4/Shift+4:YZ平面",
            "5/Shift+5:YW平面",
            "6/Shift+6:ZW平面",
            "",
            "=== 移动(WASD+QE+ZC) ===",
            "W/S:X轴 (前/后)",
            "A/D:Y轴 (左/右)",
            "Q/E:Z轴 (上/下)",
            "Z/C:W轴 (Ana/Kata)",
            f"位置 X:{tesseract.position_4d[0]:.2f}",
            f"     Y:{tesseract.position_4d[1]:.2f}",
            f"     Z:{tesseract.position_4d[2]:.2f}",
            f"     W:{tesseract.position_4d[3]:.2f}",
            "",
            "=== 切面 ===",
            f"缩放:{self.slice_scale:.2f}",
            "+/-:整体调整",
            "滚轮:单独W切面",
            "",
            "I:隐藏面板 | 0:重置"
        ]
        
        for line in controls:
            if line.startswith("==="):
                color = (200, 220, 255)
            else:
                color = (180, 180, 180)
            text = self.font_normal.render(line, True, color)
            surface.blit(text, (x, y))
            y += 18

class Visualizer:
    """主可视化器"""
    
    def __init__(self):
        self.screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
        pygame.display.set_caption("四维超立方体可视化系统")
        self.clock = pygame.time.Clock()
        
        self.tesseract = Tesseract()
        self.control_panel = ControlPanel()
        
        # 创建8个观察立方体
        self.cubes = []
        w_values = [-2.0, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, 2.0]
        positions = [
            (100, 100), (400, 100), (700, 100), (1000, 100),
            (100, 450), (400, 450), (700, 450), (1000, 450)
        ]
        for w, pos in zip(w_values, positions):
            self.cubes.append(ObservationCube(w, pos))
        
        self.auto_rotate = True
        self.auto_move = False
        self.paused = False
        self.show_info = True
        self.rotation_angles = [0.0] * 6
        self.move_time = 0.0
        self.hovered_cube = None
    
    def handle_input(self):
        """处理输入"""
        keys = pygame.key.get_pressed()
        shift_pressed = keys[K_LSHIFT] or keys[K_RSHIFT]
        
        # 手动旋转控制 - 使用数字键1-6，Shift反向
        if not self.paused and not self.auto_rotate:
            speed = self.control_panel.rotation_speed * 0.01
            direction = -1 if shift_pressed else 1
            
            if keys[K_1]: self.rotation_angles[0] += speed * direction  # XY
            if keys[K_2]: self.rotation_angles[1] += speed * direction  # XZ
            if keys[K_3]: self.rotation_angles[2] += speed * direction  # XW
            if keys[K_4]: self.rotation_angles[3] += speed * direction  # YZ
            if keys[K_5]: self.rotation_angles[4] += speed * direction  # YW
            if keys[K_6]: self.rotation_angles[5] += speed * direction  # ZW
        
        # 手动四维平移控制 - 使用WASD+QE+ZC
        if not self.auto_move and not self.paused:
            speed = self.control_panel.translation_speed
            
            # X轴控制 (W前/S后)
            if keys[K_w]: self.tesseract.position_4d[0] += speed
            if keys[K_s]: self.tesseract.position_4d[0] -= speed
            
            # Y轴控制 (A左/D右)
            if keys[K_a]: self.tesseract.position_4d[1] -= speed
            if keys[K_d]: self.tesseract.position_4d[1] += speed
            
            # Z轴控制 (Q上/E下)
            if keys[K_q]: self.tesseract.position_4d[2] += speed
            if keys[K_e]: self.tesseract.position_4d[2] -= speed
            
            # W轴控制 (Z ana/C kata，四维术语)
            if keys[K_z] and not keys[K_LCTRL]: self.tesseract.position_4d[3] += speed
            if keys[K_c]: self.tesseract.position_4d[3] -= speed
        
        # 自动旋转
        if self.auto_rotate and not self.paused:
            self.rotation_angles[0] += 0.003
            self.rotation_angles[2] += 0.005
            self.rotation_angles[5] += 0.002
        
        # 自动移动
        if self.auto_move and not self.paused:
            self.move_time += 0.05
            self.tesseract.position_4d[0] = np.sin(self.move_time * 0.3) * 1.5
            self.tesseract.position_4d[1] = np.cos(self.move_time * 0.4) * 1.2
            self.tesseract.position_4d[2] = np.sin(self.move_time * 0.5) * 1.0
            self.tesseract.position_4d[3] = np.cos(self.move_time * 0.2) * 0.8
        
        self.tesseract.update_rotation(self.rotation_angles)
        
        # 更新鼠标悬停状态
        mouse_pos = pygame.mouse.get_pos()
        self.hovered_cube = None
        for cube in self.cubes:
            if cube.contains_point(mouse_pos):
                self.hovered_cube = cube
                break
    
    def handle_events(self):
        """处理事件"""
        for event in pygame.event.get():
            if event.type == QUIT:
                return False
            elif event.type == KEYDOWN:
                if event.key == K_ESCAPE:
                    return False
                elif event.key == K_SPACE:
                    self.paused = not self.paused
                elif event.key == K_r:
                    self.auto_rotate = not self.auto_rotate
                elif event.key == K_TAB:
                    self.auto_move = not self.auto_move
                    if self.auto_move:
                        self.move_time = 0.0
                elif event.key == K_i:
                    self.show_info = not self.show_info
                elif event.key == K_0:
                    self.control_panel.slice_scale = 1.0
                    for cube in self.cubes:
                        cube.slice_scale = 1.0
                    self.rotation_angles = [0.0] * 6
                    self.tesseract.position_4d = np.array([0.0, 0.0, 0.0, 0.0])
                elif event.key == K_PLUS or event.key == K_EQUALS:
                    self.control_panel.slice_scale += 0.1
                    for cube in self.cubes:
                        cube.slice_scale = self.control_panel.slice_scale
                elif event.key == K_MINUS:
                    self.control_panel.slice_scale = max(0.1, self.control_panel.slice_scale - 0.1)
                    for cube in self.cubes:
                        cube.slice_scale = self.control_panel.slice_scale
            elif event.type == MOUSEWHEEL:
                # 滚轮控制鼠标悬停的立方体的W切面
                if self.hovered_cube:
                    delta = event.y * 0.1
                    self.hovered_cube.w_value += delta
                else:
                    # 如果没有悬停，则控制所有立方体
                    delta = event.y * 0.1
                    for cube in self.cubes:
                        cube.w_value += delta
        return True
    
    def render(self):
        """渲染场景"""
        self.screen.fill(BG_COLOR)
        
        # 渲染所有观察立方体
        for cube in self.cubes:
            surface = pygame.Surface((cube.size, cube.size), pygame.SRCALPHA)
            slice_v, slice_e = self.tesseract.get_3d_slice(cube.w_value)
            highlight = (cube == self.hovered_cube)
            cube.render(surface, slice_v, slice_e, highlight)
            self.screen.blit(surface, cube.position)
        
        # 渲染控制面板
        if self.show_info:
            panel_surface = pygame.Surface((200, SCREEN_HEIGHT), pygame.SRCALPHA)
            panel_surface.fill((0, 0, 0, 180))
            self.control_panel.render(panel_surface, self.tesseract, self.auto_rotate, self.paused, self.auto_move)
            self.screen.blit(panel_surface, (SCREEN_WIDTH - 200, 0))
        
        pygame.display.flip()
    
    def run(self):
        """主循环"""
        running = True
        while running:
            running = self.handle_events()
            self.handle_input()
            self.render()
            self.clock.tick(60)
        pygame.quit()

if __name__ == "__main__":
    visualizer = Visualizer()
    visualizer.run()


